package com.formula.datastructure.exam.graph.mst;

import com.formula.datastructure.exercise.graph.undirect.ListEdge;
import com.formula.datastructure.exercise.graph.undirect.ListGraph;
import com.formula.datastructure.exercise.graph.undirect.MatrixGraph;

public class Prim {

    private static final int INF = 99999;

    public void getMST(ListGraph graph) {
        int N = graph.vertexSize();
        ListEdge[] edges = new ListEdge[N - 1];
        int[] visited = new int[N];

        visited[0] = 1;

        ListEdge[] shortEdges = new ListEdge[N];
        for (int i = 0; i < N; i++) {
            shortEdges[i] = new ListEdge(-1, INF);
        }

        ListEdge tempEdge = graph.vertexList[0].firstEdge;
        while (tempEdge != null) {
            // 记录到这个顶点最短的边
            shortEdges[tempEdge.toIdx] = tempEdge;
            tempEdge = tempEdge.nextEdge;
        }

        //循环N-1次
        for (int k = 0; k < N - 1; k++) {
            int min = INF;
            int nextVer = 0;
            ListEdge nextEdge = null;
            //找最短边
            for (int i = 0; i < N; i++) {
                if (visited[shortEdges[i].toIdx] == 0 && min > shortEdges[i].weight) {
                    min = shortEdges[i].weight;
                    nextEdge = shortEdges[i];
                    nextVer = shortEdges[i].toIdx;
                }
            }
            edges[k] = nextEdge;
            visited[nextVer] = 1;

            //新加入的点加入新边
            tempEdge = graph.vertexList[nextVer].firstEdge;
            while (tempEdge != null) {
                if (visited[tempEdge.toIdx] == 0 && shortEdges[tempEdge.toIdx].weight > tempEdge.weight) {
                    shortEdges[tempEdge.toIdx] = tempEdge;
                }
                tempEdge = tempEdge.nextEdge;
            }
        }

    }

    public void getMST(MatrixGraph graph) {
        int N = graph.size;

        int[] visited = new int[N];
        int[] shortDis = new int[N];

        visited[0] = 1;

        for (int i = 0; i < N; i++) {
            shortDis[i] = graph.edges[0][i];
        }

        for (int i = 0; i < N; i++) {
            int min = INF;
            int nextVer = 0;
            for (int j = 0; j < N; j++) {
                if (visited[j] == 0 && shortDis[j] < min) {
                    min = shortDis[j];
                    nextVer = j;
                }
            }

            visited[nextVer] = 1;

            for (int j = 0; j < N; j++) {
                if (visited[j] == 0 && shortDis[j] > graph.edges[nextVer][j]) {
                    shortDis[j] = graph.edges[nextVer][j];
                }
            }
        }

    }
}
